{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text-to-SQL\n",
    "\n",
    "The `Text2SQL` module is a demonstration of how to use Guardrails to build a text-to-SQL workflow. It implements necessary components to build a text-to-SQL workflow, including:\n",
    "\n",
    "- [x] Connects to your database\n",
    "- [x] Given a DB schema, supports SQL validation and error handling\n",
    "- [x] Finds relevant few shot examples of `text2SQL` history, and uses them to generate a prompt\n",
    "- [x] Optionally adds modules to guard against unsafe SQL (allowing INSERT, DROP, etc.)\n",
    "- [x] Supports reasking for all invalid SQL (SQL with syntax errors, SQL targeting non-existent tables, SQL with prohibited keywords, etc.)\n",
    "\n",
    "!!! info\n",
    "    The `text2SQL` module is intended to showcase what an end-to-end `text2SQL` workflow looks like with Guardrails. As a developer, you can either use this module as a starting point for your own `text2SQL` workflow or you can use it as a reference for how to use Guardrails to build your own custom `text2SQL` workflow."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🎱  Why use Guardrails for text-to-SQL?\n",
    "\n",
    "1. **`bug-free-sql` validators:** Guardrails provides a `bug-free-sql` validator that creates a sandboxed environment to run the generated SQL against your database and check for bugs and errors.\n",
    "2. **Protection against unsafe SQL:** Guardrails allows developers to constrain the generated SQL to a subset of SQL that is safe to run against your database. For example, you can configure Guardrails to only allow `SELECT` statements, and disallow `INSERT`, `DROP`, etc.\n",
    "3. **Reasking:** Guardrails allows you to configure reasking logic, so that if the generated SQL is invalid or has bugs, you can reask the LLM to generate a new SQL. Guardrails automatically handles the reasking logic for you."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🚀 Quickstart\n",
    "\n",
    "To get started, you'll need the following:\n",
    "1. Either a connection string to your database or a database schema\n",
    "2. (Optional) Dictionary of few shot examples of `text2SQL` history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "Remove_cell"
    ]
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "from rich import print\n",
    "\n",
    "from guardrails.applications.text2sql import Text2Sql\n",
    "\n",
    "# Set your OPENAI_API_KEY as an environment variable\n",
    "# import os\n",
    "# os.environ['OPENAI_API_KEY'] = \"YOUR_API_KEY\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 📝 Setup database connection and examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"font-weight: bold\">{</span>\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'question'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'How many heads of the departments are older than 56 ?'</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'query'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'SELECT count(*) FROM head WHERE age  &gt;  56'</span>\n",
       "    <span style=\"font-weight: bold\">}</span>,\n",
       "    <span style=\"font-weight: bold\">{</span>\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'question'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'List the name, born state and age of the heads of departments ordered by age.'</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'query'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'SELECT name ,  born_state ,  age FROM head ORDER BY age'</span>\n",
       "    <span style=\"font-weight: bold\">}</span>\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m[\u001b[0m\n",
       "    \u001b[1m{\u001b[0m\n",
       "        \u001b[32m'question'\u001b[0m: \u001b[32m'How many heads of the departments are older than 56 ?'\u001b[0m,\n",
       "        \u001b[32m'query'\u001b[0m: \u001b[32m'SELECT count\u001b[0m\u001b[32m(\u001b[0m\u001b[32m*\u001b[0m\u001b[32m)\u001b[0m\u001b[32m FROM head WHERE age  >  56'\u001b[0m\n",
       "    \u001b[1m}\u001b[0m,\n",
       "    \u001b[1m{\u001b[0m\n",
       "        \u001b[32m'question'\u001b[0m: \u001b[32m'List the name, born state and age of the heads of departments ordered by age.'\u001b[0m,\n",
       "        \u001b[32m'query'\u001b[0m: \u001b[32m'SELECT name ,  born_state ,  age FROM head ORDER BY age'\u001b[0m\n",
       "    \u001b[1m}\u001b[0m\n",
       "\u001b[1m]\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "EXAMPLES = \"examples.json\"\n",
    "SQL_SCHEMA = \"schema.sql\"\n",
    "# Alternatively, you can specify a connection string instead of a schema file, like so:\n",
    "# SQL_CONN = f\"sqlite:///{os.getcwd()}department_management.sqlite\"\n",
    "\n",
    "with open(EXAMPLES, \"r\") as f:\n",
    "    examples = json.load(f)\n",
    "\n",
    "print(examples[:2])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 💡 Use Text2SQL Application\n",
    "\n",
    "- Sets up sandboxed DB based on your SQL schema / connection string\n",
    "- Finds most relevant examples and inserts them into the prompt\n",
    "- Checks that the query is valid for the schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">SELECT name FROM department ORDER BY Num_Employees DESC LIMIT <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "SELECT name FROM department ORDER BY Num_Employees DESC LIMIT \u001b[1;36m1\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "app = Text2Sql(\n",
    "    \"sqlite://\",\n",
    "    schema_file=SQL_SCHEMA,\n",
    "    examples=examples,\n",
    ")\n",
    "\n",
    "# Call the application with a natural language question.\n",
    "print(app(\"What is the name of the department with the highest number of employees?\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ❎ Demonstration on an incorrect output\n",
    "\n",
    "Below is a demonstration of how Guardrails' `Text2SQL` module handles an incorrect output from the LLM. The LLM generates an SQL that has a syntax error, and Guardrails catches this error and reasks the LLM to generate a new SQL.\n",
    "\n",
    "The output is incorrect because it refers to a table (`departments`) that does not exist in the database. Guardrails catches this error and reasks the LLM to generate a new SQL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": [
     "Remove_input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Incorrect output:\n",
       "\n",
       "<span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">\"generated_sql\"</span>: <span style=\"color: #008000; text-decoration-color: #008000\">\"SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1\"</span><span style=\"font-weight: bold\">}</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Incorrect output:\n",
       "\n",
       "\u001b[1m{\u001b[0m\u001b[32m\"generated_sql\"\u001b[0m: \u001b[32m\"SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1\"\u001b[0m\u001b[1m}\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "incorrect_llm_output = '{\"generated_sql\": \"SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1\"}'\n",
    "\n",
    "print(f\"Incorrect output:\\n\\n{incorrect_llm_output}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we can see the final corrected output, as well as the underlying logs to see what happened step by step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": [
     "Remove_input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Correct output:\n",
       "\n",
       "SELECT Name FROM department ORDER BY Num_Employees DESC LIMIT <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Correct output:\n",
       "\n",
       "SELECT Name FROM department ORDER BY Num_Employees DESC LIMIT \u001b[1;36m1\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from guardrails.llm_providers import openai_wrapper\n",
    "\n",
    "output = app.guard.parse(\n",
    "    llm_output=incorrect_llm_output,\n",
    "    llm_api=openai_wrapper,\n",
    "    engine=\"text-davinci-003\",\n",
    "    prompt_params={\n",
    "        \"nl_instruction\": \"What is the name of the department with the highest number of employees?\",\n",
    "        \"db_info\": str(app.sql_schema),\n",
    "        \"examples\": \"\",\n",
    "    },\n",
    "    max_tokens=512,\n",
    ")[\"generated_sql\"]\n",
    "\n",
    "print(f\"Correct output:\\n\\n{output}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": [
     "Remove_input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Logs\n",
       "├── ╭────────────────────────────────────────────────── Step 0 ───────────────────────────────────────────────────╮\n",
       "│   │ <span style=\"background-color: #f0f8ff\">╭──────────────────────────────────────────────── Prompt ─────────────────────────────────────────────────╮</span> │\n",
       "│   │ <span style=\"background-color: #f0f8ff\">│ No prompt                                                                                               │</span> │\n",
       "│   │ <span style=\"background-color: #f0f8ff\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span> │\n",
       "│   │ <span style=\"background-color: #f5f5dc\">╭──────────────────────────────────────────── Raw LLM Output ─────────────────────────────────────────────╮</span> │\n",
       "│   │ <span style=\"background-color: #f5f5dc\">│ {\"generated_sql\": \"SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1\"}                   │</span> │\n",
       "│   │ <span style=\"background-color: #f5f5dc\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">╭─────────────────────────────────────────── Validated Output ────────────────────────────────────────────╮</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│ {                                                                                                       │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│     'generated_sql': ReAsk(                                                                             │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│         incorrect_value='SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1',             │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│         error_message='(sqlite3.OperationalError) no such table: departments\\n[SQL: SELECT name FROM    │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│ departments ORDER BY num_employees DESC LIMIT 1]\\n(Background on this error at:                         │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│ https://sqlalche.me/e/20/e3q8)',                                                                        │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│         fix_value=None,                                                                                 │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│         path=['generated_sql']                                                                          │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│     )                                                                                                   │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">│ }                                                                                                       │</span> │\n",
       "│   │ <span style=\"background-color: #f0fff0\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span> │\n",
       "│   ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n",
       "└── ╭────────────────────────────────────────────────── Step 1 ───────────────────────────────────────────────────╮\n",
       "    │ <span style=\"background-color: #f0f8ff\">╭──────────────────────────────────────────────── Prompt ─────────────────────────────────────────────────╮</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ You are a data scientist whose job is to write SQL queries.                                             │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ Given below is XML that describes the information to extract from this document and the tags to extract │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ it into.                                                                                                │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ &lt;output&gt;                                                                                                │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     &lt;sql name=\"generated_sql\" description=\"Generate SQL for the given natural language instruction.\"    │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ format=\"bug-free-sql\"/&gt;                                                                                 │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ &lt;/output&gt;                                                                                               │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ ONLY return a valid JSON object (no other text is necessary), where the key of the field in JSON is the │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ `name` attribute of the corresponding XML, and the value is of the type specified by the corresponding  │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ XML's tag. The JSON MUST conform to the XML format, including any types and format requests e.g.        │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ requests for lists, objects and specific types. Be correct and concise.                                 │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ Here are examples of simple (XML, JSON) pairs that show the expected behavior:                          │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ - `&lt;string name='foo' format='two-words lower-case' /&gt;` =&gt; `{'foo': 'example one'}`                     │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ - `&lt;list name='bar'&gt;&lt;string format='upper-case' /&gt;&lt;/list&gt;` =&gt; `{\"bar\": ['STRING ONE', 'STRING TWO',     │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ etc.]}`                                                                                                 │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ - `&lt;object name='baz'&gt;&lt;string name=\"foo\" format=\"capitalize two-words\" /&gt;&lt;integer name=\"index\"          │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ format=\"1-indexed\" /&gt;&lt;/object&gt;` =&gt; `{'baz': {'foo': 'Some String', 'index': 1}}`                        │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ Here's schema about the database that you can use to generate the SQL query.                            │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ Try to avoid using joins if the data can be retrieved from the same table.                              │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ Table: department                                                                                       │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: Department_ID                                                                               │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: INTEGER                                                                                   │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: Name                                                                                        │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: TEXT                                                                                      │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: Creation                                                                                    │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: TEXT                                                                                      │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: Ranking                                                                                     │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: INTEGER                                                                                   │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: Budget_in_Billions                                                                          │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: REAL                                                                                      │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: Num_Employees                                                                               │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: REAL                                                                                      │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ Table: head                                                                                             │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: head_ID                                                                                     │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: INTEGER                                                                                   │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: name                                                                                        │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: TEXT                                                                                      │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: born_state                                                                                  │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: TEXT                                                                                      │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: age                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: REAL                                                                                      │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ Table: management                                                                                       │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: department_ID                                                                               │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: INTEGER                                                                                   │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         foreign_key: {'table': 'department', 'column': 'Department_ID'}                                 │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: head_ID                                                                                     │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: INTEGER                                                                                   │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         foreign_key: {'table': 'head', 'column': 'head_ID'}                                             │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     Column: temporary_acting                                                                            │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│         type: TEXT                                                                                      │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ I will give you a list of examples.                                                                     │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ I want to create a query for the following instruction:                                                 │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ What is the name of the department with the highest number of employees?                                │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ For this instruction, I was given the following JSON, which has some incorrect values.                  │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ {                                                                                                       │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│   \"generated_sql\": {                                                                                    │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     \"incorrect_value\": \"SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1\",              │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│     \"error_message\": \"(sqlite3.OperationalError) no such table: departments\\n[SQL: SELECT name FROM     │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ departments ORDER BY num_employees DESC LIMIT 1]\\n(Background on this error at:                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ https://sqlalche.me/e/20/e3q8)\"                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│   }                                                                                                     │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ }                                                                                                       │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│ Help me correct the incorrect values based on the given error messages.                                 │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">│                                                                                                         │</span> │\n",
       "    │ <span style=\"background-color: #f0f8ff\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span> │\n",
       "    │ <span style=\"background-color: #f5f5dc\">╭──────────────────────────────────────────── Raw LLM Output ─────────────────────────────────────────────╮</span> │\n",
       "    │ <span style=\"background-color: #f5f5dc\">│ {\"generated_sql\": \"SELECT Name FROM department ORDER BY Num_Employees DESC LIMIT 1\"}                    │</span> │\n",
       "    │ <span style=\"background-color: #f5f5dc\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span> │\n",
       "    │ <span style=\"background-color: #f0fff0\">╭─────────────────────────────────────────── Validated Output ────────────────────────────────────────────╮</span> │\n",
       "    │ <span style=\"background-color: #f0fff0\">│ {                                                                                                       │</span> │\n",
       "    │ <span style=\"background-color: #f0fff0\">│     'generated_sql': 'SELECT Name FROM department ORDER BY Num_Employees DESC LIMIT 1'                  │</span> │\n",
       "    │ <span style=\"background-color: #f0fff0\">│ }                                                                                                       │</span> │\n",
       "    │ <span style=\"background-color: #f0fff0\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span> │\n",
       "    ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Logs\n",
       "├── ╭────────────────────────────────────────────────── Step 0 ───────────────────────────────────────────────────╮\n",
       "│   │ \u001b[48;2;240;248;255m╭─\u001b[0m\u001b[48;2;240;248;255m───────────────────────────────────────────────\u001b[0m\u001b[48;2;240;248;255m Prompt \u001b[0m\u001b[48;2;240;248;255m────────────────────────────────────────────────\u001b[0m\u001b[48;2;240;248;255m─╮\u001b[0m │\n",
       "│   │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mNo prompt\u001b[0m\u001b[48;2;240;248;255m                                                                                              \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;248;255m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m │\n",
       "│   │ \u001b[48;2;245;245;220m╭─\u001b[0m\u001b[48;2;245;245;220m───────────────────────────────────────────\u001b[0m\u001b[48;2;245;245;220m Raw LLM Output \u001b[0m\u001b[48;2;245;245;220m────────────────────────────────────────────\u001b[0m\u001b[48;2;245;245;220m─╮\u001b[0m │\n",
       "│   │ \u001b[48;2;245;245;220m│\u001b[0m\u001b[48;2;245;245;220m \u001b[0m\u001b[48;2;245;245;220m{\"generated_sql\": \"SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1\"}\u001b[0m\u001b[48;2;245;245;220m                  \u001b[0m\u001b[48;2;245;245;220m \u001b[0m\u001b[48;2;245;245;220m│\u001b[0m │\n",
       "│   │ \u001b[48;2;245;245;220m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m╭─\u001b[0m\u001b[48;2;240;255;240m──────────────────────────────────────────\u001b[0m\u001b[48;2;240;255;240m Validated Output \u001b[0m\u001b[48;2;240;255;240m───────────────────────────────────────────\u001b[0m\u001b[48;2;240;255;240m─╮\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m{\u001b[0m\u001b[48;2;240;255;240m                                                                                                      \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m    'generated_sql': ReAsk(\u001b[0m\u001b[48;2;240;255;240m                                                                            \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m        incorrect_value='SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1',\u001b[0m\u001b[48;2;240;255;240m            \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m        error_message='(sqlite3.OperationalError) no such table: departments\\n[SQL: SELECT name FROM \u001b[0m\u001b[48;2;240;255;240m  \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240mdepartments ORDER BY num_employees DESC LIMIT 1]\\n(Background on this error at: \u001b[0m\u001b[48;2;240;255;240m                       \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240mhttps://sqlalche.me/e/20/e3q8)',\u001b[0m\u001b[48;2;240;255;240m                                                                       \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m        fix_value=None,\u001b[0m\u001b[48;2;240;255;240m                                                                                \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m        path=['generated_sql']\u001b[0m\u001b[48;2;240;255;240m                                                                         \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m    )\u001b[0m\u001b[48;2;240;255;240m                                                                                                  \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m}\u001b[0m\u001b[48;2;240;255;240m                                                                                                      \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "│   │ \u001b[48;2;240;255;240m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m │\n",
       "│   ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n",
       "└── ╭────────────────────────────────────────────────── Step 1 ───────────────────────────────────────────────────╮\n",
       "    │ \u001b[48;2;240;248;255m╭─\u001b[0m\u001b[48;2;240;248;255m───────────────────────────────────────────────\u001b[0m\u001b[48;2;240;248;255m Prompt \u001b[0m\u001b[48;2;240;248;255m────────────────────────────────────────────────\u001b[0m\u001b[48;2;240;248;255m─╮\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mYou are a data scientist whose job is to write SQL queries.\u001b[0m\u001b[48;2;240;248;255m                                            \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mGiven below is XML that describes the information to extract from this document and the tags to extract\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mit into.\u001b[0m\u001b[48;2;240;248;255m                                                                                               \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m<output>\u001b[0m\u001b[48;2;240;248;255m                                                                                               \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    <sql name=\"generated_sql\" description=\"Generate SQL for the given natural language instruction.\" \u001b[0m\u001b[48;2;240;248;255m  \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mformat=\"bug-free-sql\"/>\u001b[0m\u001b[48;2;240;248;255m                                                                                \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m</output>\u001b[0m\u001b[48;2;240;248;255m                                                                                              \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mONLY return a valid JSON object (no other text is necessary), where the key of the field in JSON is the\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m`name` attribute of the corresponding XML, and the value is of the type specified by the corresponding \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mXML's tag. The JSON MUST conform to the XML format, including any types and format requests e.g. \u001b[0m\u001b[48;2;240;248;255m      \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mrequests for lists, objects and specific types. Be correct and concise.\u001b[0m\u001b[48;2;240;248;255m                                \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mHere are examples of simple (XML, JSON) pairs that show the expected behavior:\u001b[0m\u001b[48;2;240;248;255m                         \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m- `<string name='foo' format='two-words lower-case' />` => `{'foo': 'example one'}`\u001b[0m\u001b[48;2;240;248;255m                    \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m- `<list name='bar'><string format='upper-case' /></list>` => `{\"bar\": ['STRING ONE', 'STRING TWO', \u001b[0m\u001b[48;2;240;248;255m   \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255metc.]}`\u001b[0m\u001b[48;2;240;248;255m                                                                                                \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m- `<object name='baz'><string name=\"foo\" format=\"capitalize two-words\" /><integer name=\"index\" \u001b[0m\u001b[48;2;240;248;255m        \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mformat=\"1-indexed\" /></object>` => `{'baz': {'foo': 'Some String', 'index': 1}}`\u001b[0m\u001b[48;2;240;248;255m                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mHere's schema about the database that you can use to generate the SQL query.\u001b[0m\u001b[48;2;240;248;255m                           \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mTry to avoid using joins if the data can be retrieved from the same table.\u001b[0m\u001b[48;2;240;248;255m                             \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mTable: department\u001b[0m\u001b[48;2;240;248;255m                                                                                      \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: Department_ID\u001b[0m\u001b[48;2;240;248;255m                                                                              \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: INTEGER\u001b[0m\u001b[48;2;240;248;255m                                                                                  \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: Name\u001b[0m\u001b[48;2;240;248;255m                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: TEXT\u001b[0m\u001b[48;2;240;248;255m                                                                                     \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: Creation\u001b[0m\u001b[48;2;240;248;255m                                                                                   \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: TEXT\u001b[0m\u001b[48;2;240;248;255m                                                                                     \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: Ranking\u001b[0m\u001b[48;2;240;248;255m                                                                                    \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: INTEGER\u001b[0m\u001b[48;2;240;248;255m                                                                                  \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: Budget_in_Billions\u001b[0m\u001b[48;2;240;248;255m                                                                         \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: REAL\u001b[0m\u001b[48;2;240;248;255m                                                                                     \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: Num_Employees\u001b[0m\u001b[48;2;240;248;255m                                                                              \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: REAL\u001b[0m\u001b[48;2;240;248;255m                                                                                     \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mTable: head\u001b[0m\u001b[48;2;240;248;255m                                                                                            \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: head_ID\u001b[0m\u001b[48;2;240;248;255m                                                                                    \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: INTEGER\u001b[0m\u001b[48;2;240;248;255m                                                                                  \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: name\u001b[0m\u001b[48;2;240;248;255m                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: TEXT\u001b[0m\u001b[48;2;240;248;255m                                                                                     \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: born_state\u001b[0m\u001b[48;2;240;248;255m                                                                                 \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: TEXT\u001b[0m\u001b[48;2;240;248;255m                                                                                     \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: age\u001b[0m\u001b[48;2;240;248;255m                                                                                        \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: REAL\u001b[0m\u001b[48;2;240;248;255m                                                                                     \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mTable: management\u001b[0m\u001b[48;2;240;248;255m                                                                                      \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: department_ID\u001b[0m\u001b[48;2;240;248;255m                                                                              \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: INTEGER\u001b[0m\u001b[48;2;240;248;255m                                                                                  \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        foreign_key: {'table': 'department', 'column': 'Department_ID'}\u001b[0m\u001b[48;2;240;248;255m                                \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: head_ID\u001b[0m\u001b[48;2;240;248;255m                                                                                    \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: INTEGER\u001b[0m\u001b[48;2;240;248;255m                                                                                  \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        foreign_key: {'table': 'head', 'column': 'head_ID'}\u001b[0m\u001b[48;2;240;248;255m                                            \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    Column: temporary_acting\u001b[0m\u001b[48;2;240;248;255m                                                                           \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m        type: TEXT\u001b[0m\u001b[48;2;240;248;255m                                                                                     \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mI will give you a list of examples.\u001b[0m\u001b[48;2;240;248;255m                                                                    \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mI want to create a query for the following instruction:\u001b[0m\u001b[48;2;240;248;255m                                                \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mWhat is the name of the department with the highest number of employees?\u001b[0m\u001b[48;2;240;248;255m                               \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mFor this instruction, I was given the following JSON, which has some incorrect values.\u001b[0m\u001b[48;2;240;248;255m                 \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m{\u001b[0m\u001b[48;2;240;248;255m                                                                                                      \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m  \"generated_sql\": {\u001b[0m\u001b[48;2;240;248;255m                                                                                   \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    \"incorrect_value\": \"SELECT name FROM departments ORDER BY num_employees DESC LIMIT 1\",\u001b[0m\u001b[48;2;240;248;255m             \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m    \"error_message\": \"(sqlite3.OperationalError) no such table: departments\\n[SQL: SELECT name FROM \u001b[0m\u001b[48;2;240;248;255m   \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mdepartments ORDER BY num_employees DESC LIMIT 1]\\n(Background on this error at: \u001b[0m\u001b[48;2;240;248;255m                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mhttps://sqlalche.me/e/20/e3q8)\"\u001b[0m\u001b[48;2;240;248;255m                                                                        \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m  }\u001b[0m\u001b[48;2;240;248;255m                                                                                                    \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m}\u001b[0m\u001b[48;2;240;248;255m                                                                                                      \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255mHelp me correct the incorrect values based on the given error messages.\u001b[0m\u001b[48;2;240;248;255m                                \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m│\u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m                                                                                                       \u001b[0m\u001b[48;2;240;248;255m \u001b[0m\u001b[48;2;240;248;255m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;248;255m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m │\n",
       "    │ \u001b[48;2;245;245;220m╭─\u001b[0m\u001b[48;2;245;245;220m───────────────────────────────────────────\u001b[0m\u001b[48;2;245;245;220m Raw LLM Output \u001b[0m\u001b[48;2;245;245;220m────────────────────────────────────────────\u001b[0m\u001b[48;2;245;245;220m─╮\u001b[0m │\n",
       "    │ \u001b[48;2;245;245;220m│\u001b[0m\u001b[48;2;245;245;220m \u001b[0m\u001b[48;2;245;245;220m{\"generated_sql\": \"SELECT Name FROM department ORDER BY Num_Employees DESC LIMIT 1\"}\u001b[0m\u001b[48;2;245;245;220m                   \u001b[0m\u001b[48;2;245;245;220m \u001b[0m\u001b[48;2;245;245;220m│\u001b[0m │\n",
       "    │ \u001b[48;2;245;245;220m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m │\n",
       "    │ \u001b[48;2;240;255;240m╭─\u001b[0m\u001b[48;2;240;255;240m──────────────────────────────────────────\u001b[0m\u001b[48;2;240;255;240m Validated Output \u001b[0m\u001b[48;2;240;255;240m───────────────────────────────────────────\u001b[0m\u001b[48;2;240;255;240m─╮\u001b[0m │\n",
       "    │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m{\u001b[0m\u001b[48;2;240;255;240m                                                                                                      \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m    'generated_sql': 'SELECT Name FROM department ORDER BY Num_Employees DESC LIMIT 1'\u001b[0m\u001b[48;2;240;255;240m                 \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;255;240m│\u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m}\u001b[0m\u001b[48;2;240;255;240m                                                                                                      \u001b[0m\u001b[48;2;240;255;240m \u001b[0m\u001b[48;2;240;255;240m│\u001b[0m │\n",
       "    │ \u001b[48;2;240;255;240m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m │\n",
       "    ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "app.guard.history.last.tree"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ⚒️ Customizing the `Text2SQL` module\n",
    "\n",
    "You can customize the `Text2SQL` module to fit your needs. Here are some examples of how you can customize the module:\n",
    "\n",
    "1. **Prompt:** You can customize the `text2SQL` prompt by passing in a custom `RAIL` spec to the module in the `rail_spec` argument.\n",
    "2. **Few shot examples:** You can customize how the few shot examples are stored, retrieved, and inserted into the prompt by updating the arguments `vector_db`, `document_store`, `num_relevant_examples` and `example_formatter`.\n",
    "3. **Reask prompt:** You can customize the reask prompt by updating the `reask_prompt` argument.\n",
    "4. **LLM API:** You can customize the LLM API by updating the `llm_api` argument. To pass in specific kwargs to the LLM API, you can update the `llm_api_kwargs` argument (e.g. `llm_api_kwargs={'max_length': 100}`)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tiff-env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
